import json
import argparse
import os
import open3d as o3d
import numpy as np
import math
from my_io import mkdir, copy_file
from tqdm import tqdm
from pandaset_devkit.python.pandaset import geometry
from pandaset_devkit.python.pandaset import DataSet
from utils import bbox3d2corners, group_rectangle_vertexs, group_plane_equation, points_in_bboxes

valid_labels = ['Car', 'Pedestrian', 'Truck', 'Motorcycle']

def swap_xy(pc_ori):
    pc = pc_ori.astype(np.float32)
    pc[:, 0] = pc_ori[:, 1]
    pc[:, 1] = -pc_ori[:, 0]
    return pc
def quart_to_rpy(x, y, z, w):
    roll = math.atan2(2 * (w * x + y * z), 1 - 2 * (x * x + y * y))
    pitch = math.asin(2 * (w * y - x * z))
    yaw = math.atan2(2 * (w * z + x * y), 1 - 2 * (z * z + y * y))
    return roll, pitch, yaw

def get_valid_bboxs_with_enough_pts(sampled_bboxes, ego_points, label_list):
    if len(sampled_bboxes) == 0:
        return []
    bboxes_corners = bbox3d2corners(np.stack(sampled_bboxes, axis=0))  # (n, 8, 3) 3dbox转8个顶点
    bbox_group_rectangle_vertexs = group_rectangle_vertexs(bboxes_corners)  # (n, 6, 4, 3) 8个顶点转6个面
    group_plane_equation_params = group_plane_equation(bbox_group_rectangle_vertexs)  # 4顶点平面转平面方程
    masks = points_in_bboxes(ego_points, group_plane_equation_params)  # (N, n)
    valid_label_list = []
    for id_mask in range(masks.shape[1]):
        valid_pt_size = np.sum(masks[:, id_mask] == True)
        if valid_pt_size > 50:
            valid_label_list.append(label_list[id_mask])
    return valid_label_list

def main(args):
    assert os.path.exists(args.data_path)
    pandaset = DataSet(args.data_path)
    datasets = os.listdir(args.data_path)
    datasets.sort(key=lambda x: x)
    for seq in tqdm(datasets):
        if args.start_idx != -1 and args.start_idx > int(seq[-3:]):
            continue
        print("开始转换， seq:", seq)
        dataset_path = os.path.join(args.data_path, seq)
        main_worker(pandaset[seq], dataset_path)

def main_worker(seq, dataset_path):
    seq.load_lidar().load_cuboids()
    seq.load_semseg()
    if seq.semseg == None:
        have_seq = False
    else:
        have_seq = True

    cuboids_pkl_path = os.path.join(dataset_path, "annotations", "cuboids")
    cuboids_txt_path = os.path.join(dataset_path, "annotations_txt", "cuboids", "all")
    cuboids_pandar64_txt_path = os.path.join(dataset_path, "annotations_txt", "cuboids", "pandar64")
    cuboids_pandarGT_txt_path = os.path.join(dataset_path, "annotations_txt", "cuboids", "pandarGT")
    if have_seq:
        semseg_pandar64_txt_path = os.path.join(dataset_path, "annotations_txt", "semseg", "pandar64")
        semseg_pandarGT_txt_path = os.path.join(dataset_path, "annotations_txt", "semseg", "pandarGT")
    lidar_pandar64_bin_paths = os.path.join(dataset_path, "lidar_bin", "pandar64")
    lidar_pandarGT_bin_paths = os.path.join(dataset_path, "lidar_bin", "pandarGT")
    # radar_paths = os.path.join(dataset_path, "radar")
    # radar_bin_paths = os.path.join(dataset_path, "radar_bin")

    mkdir(cuboids_txt_path)
    mkdir(cuboids_pandar64_txt_path)
    mkdir(cuboids_pandarGT_txt_path)
    if have_seq:
        mkdir(semseg_pandar64_txt_path)
        mkdir(semseg_pandarGT_txt_path)
    mkdir(lidar_pandar64_bin_paths)
    mkdir(lidar_pandarGT_bin_paths)
    # mkdir(radar_bin_paths)
    if have_seq:
        copy_file(os.path.join(dataset_path, "annotations", "semseg", "classes.json"), os.path.join(dataset_path, "annotations_txt", "semseg", "classes.json"))

    labels = os.listdir(cuboids_pkl_path)
    labels.sort(key=lambda x: x)
    for label in tqdm(labels):
        frame = str.split(label, '.')[0]
        id = int(frame)
        seq.lidar.set_sensor(-1)
        if have_seq:
            pandar64_semseg = seq.semseg[id][seq.lidar[id]['d'] == 0].to_numpy()
            pandarGT_semseg = seq.semseg[id][seq.lidar[id]['d'] == 1].to_numpy()

        seq.lidar.set_sensor(0)
        pandar64_points = seq.lidar[id].to_numpy()[..., :4]
        seq.lidar.set_sensor(1)
        pandarGT_points = seq.lidar[id].to_numpy()[..., :4]
        poses = seq.lidar.poses[id]
        _, _, rotate_yaw = quart_to_rpy(poses['heading']['x'], poses['heading']['y'], poses['heading']['z'], poses['heading']['w'])
        cuboids = seq.cuboids[id]
        ego_pandar64_points = geometry.lidar_points_to_ego(pandar64_points[:, :3], poses)
        ego_pandar64_points = swap_xy(ego_pandar64_points)
        ego_pandarGT_points = geometry.lidar_points_to_ego(pandarGT_points[:, :3], poses)
        ego_pandarGT_points = swap_xy(ego_pandarGT_points)
        ego_pandar64_points[:, 2] += 0.3
        ego_pandarGT_points[:, 2] += 0.3

        label_list_pandar64 = []
        label_list_pandarGT = []
        label_list = []
        sampled_pandar64_bboxes = []
        sampled_pandarGT_bboxes = []
        for i, row in cuboids.iterrows():
            w, l, h = row["dimensions.x"], row["dimensions.y"], row["dimensions.z"]
            yaw = row["yaw"] - rotate_yaw
            while yaw < -np.pi:
                yaw = np.pi * 2 + yaw
            while yaw > np.pi:
                yaw = yaw - np.pi * 2
            center_xyz = np.array([[row["position.x"], row["position.y"], row["position.z"]]])
            rotate_corners = geometry.lidar_points_to_ego(center_xyz, poses)
            x, y, z = rotate_corners[0, 0], rotate_corners[0, 1], rotate_corners[0, 2]
            z += 0.3

            if x < -24 or x > 72 or y < -24 or y > 24:
                continue

            label_name = row['label']
            if ' ' in label_name:
                replace = True
                for label in valid_labels:
                    if label in label_name:
                        label_name = label
                        replace = False
                if replace:
                    label_name = label_name.replace(' / ', '/')
                    label_name = label_name.replace('/', '-')
                    label_name = label_name.replace(' - ', '-')
                    label_name = label_name.replace(' ', '-')

            temp = ""
            for key in [label_name, y, -x, z, l, w, h, yaw]:
                temp += str(key)
                temp += " "
            temp += "\n"
            label_list.append(temp)
            cur_bbox = np.array([y, -x, z, l, w, h, yaw])
            if row["cuboids.sensor_id"] != 0:
                sampled_pandarGT_bboxes.append(cur_bbox)
                label_list_pandarGT.append(temp)
            if row["cuboids.sensor_id"] != 1:
                sampled_pandar64_bboxes.append(cur_bbox)
                label_list_pandar64.append(temp)

        label_list_pandar64 = get_valid_bboxs_with_enough_pts(sampled_pandar64_bboxes, ego_pandar64_points, label_list_pandar64)
        label_list_pandarGT = get_valid_bboxs_with_enough_pts(sampled_pandarGT_bboxes, ego_pandarGT_points, label_list_pandarGT)

        txt_name = os.path.join(cuboids_txt_path, frame + ".txt")
        with open(txt_name, "w") as f:
            for label in label_list:
                f.write(label)
        txt_name = os.path.join(cuboids_pandar64_txt_path, frame + ".txt")
        with open(txt_name, "w") as f:
            for label in label_list_pandar64:
                f.write(label)
        txt_name = os.path.join(cuboids_pandarGT_txt_path, frame + ".txt")
        with open(txt_name, "w") as f:
            for label in label_list_pandarGT:
                f.write(label)

        ego_pandar64_points.tofile(os.path.join(lidar_pandar64_bin_paths, frame + '.bin'))
        ego_pandarGT_points.tofile(os.path.join(lidar_pandarGT_bin_paths, frame + '.bin'))
        if have_seq:
            sem = pandar64_semseg.astype(np.float32)
            sem.tofile(os.path.join(semseg_pandar64_txt_path, frame + '.bin'))
            sem = pandarGT_semseg.astype(np.float32)
            sem.tofile(os.path.join(semseg_pandarGT_txt_path, frame + '.bin'))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Configuration Parameters')
    parser.add_argument('--start-idx', default=-1, help='start process id')
    parser.add_argument('--data-path', default='/home/adt/deeplearning/Data/pandaset/pandaset_2', help='your data root path')
    args = parser.parse_args()
    main(args)
